"""
结果输出和后处理模块

该模块提供了计算结果的:
1. 数据保存功能
2. 可视化功能
3. 数据分析功能
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from typing import Dict, List, Tuple, Optional
import json
import os
from dataclasses import dataclass

@dataclass
class ChannelResults:
    """子通道计算结果数据类"""
    channel_id: int
    axial_positions: np.ndarray
    pressures: np.ndarray
    temperatures: np.ndarray
    velocities: np.ndarray
    void_fractions: np.ndarray
    qualities: np.ndarray
    heat_fluxes: np.ndarray

class OutputProcessor:
    """结果输出和后处理类"""
    
    def __init__(self, output_dir: str):
        """
        初始化输出处理器
        
        参数:
            output_dir: 输出目录路径
        """
        self.output_dir = output_dir
        os.makedirs(output_dir, exist_ok=True)
        
    def save_results(self, results: List[ChannelResults], 
                    simulation_params: Dict) -> None:
        """
        保存计算结果
        
        参数:
            results: 子通道计算结果列表
            simulation_params: 模拟参数字典
        """
        # 保存模拟参数
        params_file = os.path.join(self.output_dir, 'simulation_params.json')
        with open(params_file, 'w') as f:
            json.dump(simulation_params, f, indent=4)
            
        # 保存计算结果
        for result in results:
            channel_dir = os.path.join(self.output_dir, 
                                     f'channel_{result.channel_id}')
            os.makedirs(channel_dir, exist_ok=True)
            
            # 创建DataFrame
            df = pd.DataFrame({
                'axial_position': result.axial_positions,
                'pressure': result.pressures,
                'temperature': result.temperatures,
                'velocity': result.velocities,
                'void_fraction': result.void_fractions,
                'quality': result.qualities,
                'heat_flux': result.heat_fluxes
            })
            
            # 保存为CSV文件
            csv_file = os.path.join(channel_dir, 'results.csv')
            df.to_csv(csv_file, index=False)
            
    def plot_axial_distributions(self, results: List[ChannelResults],
                               variable: str, title: Optional[str] = None,
                               save: bool = True) -> None:
        """
        绘制轴向分布图
        
        参数:
            results: 子通道计算结果列表
            variable: 要绘制的变量名
            title: 图标题
            save: 是否保存图片
        """
        plt.figure(figsize=(10, 6))
        
        for result in results:
            data = getattr(result, variable + 's')  # 添加's'以匹配属性名
            plt.plot(result.axial_positions, data, 
                    label=f'Channel {result.channel_id}')
            
        plt.xlabel('Axial Position (m)')
        plt.ylabel(self._get_variable_label(variable))
        plt.title(title or f'Axial Distribution of {variable.capitalize()}')
        plt.grid(True)
        plt.legend()
        
        if save:
            plt.savefig(os.path.join(self.output_dir, 
                                    f'{variable}_distribution.png'))
        plt.close()
        
    def plot_channel_comparison(self, results: List[ChannelResults],
                              variables: List[str], 
                              axial_position: float) -> None:
        """
        绘制子通道间的参数对比图
        
        参数:
            results: 子通道计算结果列表
            variables: 要对比的变量列表
            axial_position: 轴向位置
        """
        n_vars = len(variables)
        fig, axes = plt.subplots(1, n_vars, figsize=(5*n_vars, 6))
        
        if n_vars == 1:
            axes = [axes]
            
        channel_ids = [r.channel_id for r in results]
        
        for ax, var in zip(axes, variables):
            values = []
            for result in results:
                # 找到最接近指定轴向位置的索引
                idx = np.abs(result.axial_positions - axial_position).argmin()
                data = getattr(result, var + 's')
                values.append(data[idx])
                
            ax.bar(channel_ids, values)
            ax.set_xlabel('Channel ID')
            ax.set_ylabel(self._get_variable_label(var))
            ax.set_title(f'{var.capitalize()} at z = {axial_position:.3f} m')
            ax.grid(True)
            
        plt.tight_layout()
        plt.savefig(os.path.join(self.output_dir, 
                                f'channel_comparison_z{axial_position:.3f}.png'))
        plt.close()
        
    def analyze_results(self, results: List[ChannelResults]) -> Dict:
        """
        分析计算结果,生成统计报告
        
        参数:
            results: 子通道计算结果列表
            
        返回:
            包含分析结果的字典
        """
        analysis = {}
        
        # 计算各参数的最大、最小和平均值
        variables = ['pressure', 'temperature', 'velocity', 
                    'void_fraction', 'quality', 'heat_flux']
        
        for var in variables:
            var_data = []
            for result in results:
                data = getattr(result, var + 's')
                var_data.extend(data)
                
            var_data = np.array(var_data)
            analysis[var] = {
                'max': float(np.max(var_data)),
                'min': float(np.min(var_data)),
                'mean': float(np.mean(var_data)),
                'std': float(np.std(var_data))
            }
            
        # 计算子通道间的参数差异
        for var in variables:
            max_diff = 0
            for i, r1 in enumerate(results):
                for r2 in results[i+1:]:
                    data1 = getattr(r1, var + 's')
                    data2 = getattr(r2, var + 's')
                    diff = np.max(np.abs(data1 - data2))
                    max_diff = max(max_diff, diff)
                    
            analysis[f'{var}_max_channel_difference'] = float(max_diff)
            
        # 保存分析结果
        analysis_file = os.path.join(self.output_dir, 'analysis_report.json')
        with open(analysis_file, 'w') as f:
            json.dump(analysis, f, indent=4)
            
        return analysis
        
    def _get_variable_label(self, variable: str) -> str:
        """获取变量的标签"""
        labels = {
            'pressure': 'Pressure (Pa)',
            'temperature': 'Temperature (K)',
            'velocity': 'Velocity (m/s)',
            'void_fraction': 'Void Fraction',
            'quality': 'Quality',
            'heat_flux': 'Heat Flux (W/m²)'
        }
        return labels.get(variable, variable.capitalize())
        
    def generate_report(self, results: List[ChannelResults],
                       simulation_params: Dict,
                       analysis_results: Dict) -> None:
        """
        生成HTML格式的报告
        
        参数:
            results: 子通道计算结果列表
            simulation_params: 模拟参数
            analysis_results: 分析结果
        """
        # 创建报告目录
        report_dir = os.path.join(self.output_dir, 'report')
        os.makedirs(report_dir, exist_ok=True)
        
        # 复制所有图片到报告目录
        for file in os.listdir(self.output_dir):
            if file.endswith('.png'):
                src = os.path.join(self.output_dir, file)
                dst = os.path.join(report_dir, file)
                if os.path.exists(src):
                    import shutil
                    shutil.copy2(src, dst)
                    
        # 生成HTML报告
        html_content = f"""
        <!DOCTYPE html>
        <html>
        <head>
            <title>Subchannel Analysis Report</title>
            <style>
                body {{ font-family: Arial, sans-serif; margin: 40px; }}
                h1, h2 {{ color: #333; }}
                table {{ border-collapse: collapse; width: 100%; }}
                th, td {{ border: 1px solid #ddd; padding: 8px; text-align: left; }}
                th {{ background-color: #f2f2f2; }}
                img {{ max-width: 100%; height: auto; margin: 20px 0; }}
            </style>
        </head>
        <body>
            <h1>Subchannel Analysis Report</h1>
            
            <h2>Simulation Parameters</h2>
            <table>
                <tr><th>Parameter</th><th>Value</th></tr>
                {''.join(f"<tr><td>{k}</td><td>{v}</td></tr>" 
                        for k, v in simulation_params.items())}
            </table>
            
            <h2>Analysis Results</h2>
            <table>
                <tr><th>Variable</th><th>Maximum</th><th>Minimum</th>
                    <th>Mean</th><th>Std Dev</th></tr>
                {''.join(f"<tr><td>{var}</td><td>{data['max']:.3f}</td>"
                        f"<td>{data['min']:.3f}</td><td>{data['mean']:.3f}</td>"
                        f"<td>{data['std']:.3f}</td></tr>"
                        for var, data in analysis_results.items()
                        if isinstance(data, dict))}
            </table>
            
            <h2>Visualization</h2>
            {''.join(f'<img src="{file}" alt="{file}"><br>'
                    for file in os.listdir(report_dir)
                    if file.endswith('.png'))}
        </body>
        </html>
        """
        
        # 保存HTML报告
        report_file = os.path.join(report_dir, 'report.html')
        with open(report_file, 'w') as f:
            f.write(html_content) 